from sqlalchemy import Column, Integer, String
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker

from init.mysql import engine

Base = declarative_base()


class StockPreForecast(Base):
    __tablename__ = 'stock_pre_forecast'
    id = Column(Integer, primary_key=True, autoincrement=True)
    code = Column(String(10))
    pdate = Column(String(10))
    rdate = Column(String(10))
    type = Column(String(10))
    abs = Column(String(200))
    old = Column(String(20))

    def __init__(self, code, pdate, rdate, type, abs, old):
        self.code = code
        self.pdate = pdate
        self.rdate = rdate
        self.type = type
        self.abs = abs
        self.old = old

    @staticmethod
    def insert_or_update(code, data_list):
        session = sessionmaker(bind=engine)()
        for data in data_list:
            stock_pre_forecast = session.query(StockPreForecast).filter_by(code=code, pdate=data['pdate'],
                                                                          rdate=data['rdate']).first()
            if stock_pre_forecast:
                stock_pre_forecast.pdate = data['pdate']
                stock_pre_forecast.rdate = data['rdate']
                stock_pre_forecast.type = data['type']
                stock_pre_forecast.abs = data['abs']
                stock_pre_forecast.old = data['old']

            else:
                session.add(StockPreForecast(code=code, pdate=data['pdate'], rdate=data['rdate'], type=data['type'],
                                            abs=data['abs'], old=data['old']))
        session.commit()


Base.metadata.create_all(engine)
